#include "RayTracingMath.h"
#include <algorithm>
#include "Templates/NewObject.h"

bool compareX(const Triangle& t1, const Triangle& t2)
{
	return t1.center.x < t2.center.x;
}

bool compareY(const Triangle& t1, const Triangle& t2)
{
	return t1.center.y < t2.center.y;
}

bool compareZ(const Triangle& t1, const Triangle& t2)
{
	return t1.center.z < t2.center.z;
}

float HitTriangle(Triangle* triangle, Ray ray)
{
	glm::vec3 p1 = triangle->p1, p2 = triangle->p2, p3 = triangle->p3;
	glm::vec3 S = ray.mOrigin;
	glm::vec3 d = ray.mDirection;
	glm::vec3 N = glm::normalize(glm::cross(p2 - p1, p3 - p1));
	if (glm::dot(N, d) > 0.0f) N = -N;

	if (fabs(glm::dot(N, d)) < 0.00001f) return FLT_MAX;

	float t = (glm::dot(N, p1) - glm::dot(S, N)) / glm::dot(d, N);
	if (t < 0.0005f) return FLT_MAX;

	glm::vec3 P = S + d * t;

	glm::vec3 c1 = glm::cross(p2 - p1, P - p1);
	glm::vec3 c2 = glm::cross(p3 - p2, P - p2);
	glm::vec3 c3 = glm::cross(p1 - p3, P - p3);
	if (glm::dot(c1, N) > 0 && glm::dot(c2, N) > 0 && glm::dot(c3, N) > 0) return t;
	if (glm::dot(c1, N) < 0 && glm::dot(c2, N) < 0 && glm::dot(c3, N) < 0) return t;

	return FLT_MAX;
}

BVHNode* BuildBVH(std::vector<Triangle>& triangles, int l, int r, int n)
{
	if (l > r)
		return nullptr;

	BVHNode* node = NewObject<BVHNode>();

	for (int i = l; i <= r; ++i)
	{
		float minx = Math::Min(triangles[i].p1.x, Math::Min(triangles[i].p2.x, triangles[i].p3.x));
		float miny = Math::Min(triangles[i].p1.y, Math::Min(triangles[i].p2.y, triangles[i].p3.y));
		float minz = Math::Min(triangles[i].p1.z, Math::Min(triangles[i].p2.z, triangles[i].p3.z));
		node->aa.x = Math::Min(node->aa.x, minx);
		node->aa.y = Math::Min(node->aa.y, miny);
		node->aa.z = Math::Min(node->aa.z, minz);

		float maxx = Math::Max(triangles[i].p1.x, Math::Max(triangles[i].p2.x, triangles[i].p3.x));
		float maxy = Math::Max(triangles[i].p1.y, Math::Max(triangles[i].p2.y, triangles[i].p3.y));
		float maxz = Math::Max(triangles[i].p1.z, Math::Max(triangles[i].p2.z, triangles[i].p3.z));
		node->bb.x = Math::Max(node->bb.x, maxx);
		node->bb.y = Math::Max(node->bb.y, maxy);
		node->bb.z = Math::Max(node->bb.z, maxz);
	}

	if ((r - l + 1) <= n)
	{
		node->n = r - l + 1;
		node->l = l;
		return node;
	}

	float lenx = node->bb.x - node->aa.x;
	float leny = node->bb.y - node->aa.y;
	float lenz = node->bb.z - node->aa.z;

	if (lenx >= leny && lenx >= lenz)
		std::sort(triangles.begin() + l, triangles.begin() + r + 1, compareX);

	if (leny >= lenx && leny >= lenz)
		std::sort(triangles.begin() + l, triangles.begin() + r + 1, compareY);

	if (lenz >= lenx && lenz >= leny)
		std::sort(triangles.begin() + l, triangles.begin() + r + 1, compareZ);

	int mid = (l + r) / 2;
	node->left = BuildBVH(triangles, l, mid, n);
	node->right = BuildBVH(triangles, mid + 1, r, n);

	return node;
}

BVHNode* BuildBVHWithSAH(std::vector<Triangle>& triangles, int l, int r, int n)
{
	if (l > r)
		return nullptr;

	BVHNode* node = NewObject<BVHNode>();

	for (int i = l; i <= r; ++i)
	{
		float minx = Math::Min(triangles[i].p1.x, Math::Min(triangles[i].p2.x, triangles[i].p3.x));
		float miny = Math::Min(triangles[i].p1.y, Math::Min(triangles[i].p2.y, triangles[i].p3.y));
		float minz = Math::Min(triangles[i].p1.z, Math::Min(triangles[i].p2.z, triangles[i].p3.z));
		node->aa.x = Math::Min(node->aa.x, minx);
		node->aa.y = Math::Min(node->aa.y, miny);
		node->aa.z = Math::Min(node->aa.z, minz);

		float maxx = Math::Max(triangles[i].p1.x, Math::Max(triangles[i].p2.x, triangles[i].p3.x));
		float maxy = Math::Max(triangles[i].p1.y, Math::Max(triangles[i].p2.y, triangles[i].p3.y));
		float maxz = Math::Max(triangles[i].p1.z, Math::Max(triangles[i].p2.z, triangles[i].p3.z));
		node->bb.x = Math::Max(node->bb.x, maxx);
		node->bb.y = Math::Max(node->bb.y, maxy);
		node->bb.z = Math::Max(node->bb.z, maxz);
	}

	if ((r - l + 1) <= n)
	{
		node->n = r - l + 1;
		node->l = l;
		return node;
	}


	float Cost = FLT_MAX;
	int Axis = 0;
	int Split = (l + r) / 2;
	for (int axis = 0; axis < 3; axis++) {
	
		if (axis == 0) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareX);
		if (axis == 1) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareY);
		if (axis == 2) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareZ);


		std::vector<glm::vec3> leftMax(r - l + 1, glm::vec3(FLT_MIN, FLT_MIN, FLT_MIN));
		std::vector<glm::vec3> leftMin(r - l + 1, glm::vec3(FLT_MAX, FLT_MAX, FLT_MAX));

		for (int i = l; i <= r; i++) {
			Triangle& t = triangles[i];
			int bias = (i == l) ? 0 : 1; 

			leftMax[i - l].x = Math::Max(leftMax[i - l - bias].x, Math::Max(t.p1.x, Math::Max(t.p2.x, t.p3.x)));
			leftMax[i - l].y = Math::Max(leftMax[i - l - bias].y, Math::Max(t.p1.y, Math::Max(t.p2.y, t.p3.y)));
			leftMax[i - l].z = Math::Max(leftMax[i - l - bias].z, Math::Max(t.p1.z, Math::Max(t.p2.z, t.p3.z)));

			leftMin[i - l].x = Math::Min(leftMin[i - l - bias].x, Math::Min(t.p1.x, Math::Min(t.p2.x, t.p3.x)));
			leftMin[i - l].y = Math::Min(leftMin[i - l - bias].y, Math::Min(t.p1.y, Math::Min(t.p2.y, t.p3.y)));
			leftMin[i - l].z = Math::Min(leftMin[i - l - bias].z, Math::Min(t.p1.z, Math::Min(t.p2.z, t.p3.z)));
		}


		std::vector<glm::vec3> rightMax(r - l + 1, glm::vec3(FLT_MIN, FLT_MIN, FLT_MIN));
		std::vector<glm::vec3> rightMin(r - l + 1, glm::vec3(FLT_MAX, FLT_MAX, FLT_MAX));
		
		for (int i = r; i >= l; i--) {
			Triangle& t = triangles[i];
			int bias = (i == r) ? 0 : 1;

			rightMax[i - l].x = Math::Max(rightMax[i - l + bias].x, Math::Max(t.p1.x, Math::Max(t.p2.x, t.p3.x)));
			rightMax[i - l].y = Math::Max(rightMax[i - l + bias].y, Math::Max(t.p1.y, Math::Max(t.p2.y, t.p3.y)));
			rightMax[i - l].z = Math::Max(rightMax[i - l + bias].z, Math::Max(t.p1.z, Math::Max(t.p2.z, t.p3.z)));

			rightMin[i - l].x = Math::Min(rightMin[i - l + bias].x, Math::Min(t.p1.x, Math::Min(t.p2.x, t.p3.x)));
			rightMin[i - l].y = Math::Min(rightMin[i - l + bias].y, Math::Min(t.p1.y, Math::Min(t.p2.y, t.p3.y)));
			rightMin[i - l].z = Math::Min(rightMin[i - l + bias].z, Math::Min(t.p1.z, Math::Min(t.p2.z, t.p3.z)));
		}


		float cost = FLT_MAX;
		int split = l;
		for (int i = l; i <= r - 1; i++) {
			float lenx, leny, lenz;
	
			glm::vec3 leftAA = leftMin[i - l];
			glm::vec3 leftBB = leftMax[i - l];
			lenx = leftBB.x - leftAA.x;
			leny = leftBB.y - leftAA.y;
			lenz = leftBB.z - leftAA.z;
			float leftS = 2.0 * ((lenx * leny) + (lenx * lenz) + (leny * lenz));
			float leftCost = leftS * (i - l + 1);

			glm::vec3 rightAA = rightMin[i + 1 - l];
			glm::vec3 rightBB = rightMax[i + 1 - l];
			lenx = rightBB.x - rightAA.x;
			leny = rightBB.y - rightAA.y;
			lenz = rightBB.z - rightAA.z;
			float rightS = 2.0 * ((lenx * leny) + (lenx * lenz) + (leny * lenz));
			float rightCost = rightS * (r - i);


			float totalCost = leftCost + rightCost;
			if (totalCost < cost) {
				cost = totalCost;
				split = i;
			}
		}

		if (cost < Cost) {
			Cost = cost;
			Axis = axis;
			Split = split;
		}
	}

	if (Axis == 0) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareX);
	if (Axis == 1) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareY);
	if (Axis == 2) std::sort(&triangles[0] + l, &triangles[0] + r + 1, compareZ);

	node->left = BuildBVHWithSAH(triangles, l, Split, n);
	node->right = BuildBVHWithSAH(triangles, Split + 1, r, n);

	return node;
}

void FreeBvh(BVHNode* node)
{
	if (node == nullptr)
		return;

	FreeBvh(node->left);
	FreeBvh(node->right);

	DeleteObject<BVHNode>(node);
}

bool IsHitNode(Ray& ray, BVHNode* node)
{
	if (node == nullptr)
		return false;

	glm::vec3 inPoint = (node->aa - ray.mOrigin) / ray.mDirection;
	glm::vec3 outPoint = (node->bb - ray.mOrigin) / ray.mDirection;

	glm::vec3 tmax = (glm::max)(inPoint, outPoint);
	glm::vec3 tmin = (glm::min)(inPoint, outPoint);

	float t0 = (glm::max)(tmin.x, (glm::max)(tmin.y, tmin.z));
	float t1 = (glm::min)(tmax.x, (glm::min)(tmax.y, tmax.z));

	float t = (t1 >= t0) ? (t0 > 0.f ? t0 : t1) : -1;

	return t > 0;
}

HitResult HitNodeResult(std::vector<Triangle>& triangles, Ray& ray, BVHNode* node)
{
	HitResult result;

	result.node = node;
#if 0
	result.distance = glm::length((node->aa + node->bb) * 0.5f - ray.mOrigin);
#else
	for (int i = node->l; i <= (node->l + node->n - 1); ++i)
	{
		float d = HitTriangle(&triangles[i], ray);
		if (d < FLT_MAX && d < result.distance)
		{
			result.triangle = &triangles[i];
			result.distance = d;
		}
	}
#endif

	return result;
}
